{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ScitaPqhKtuW"
      },
      "source": [
        "##### Copyright 2021 The TensorFlow Hub Authors.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jvztxQ6VsK2k"
      },
      "outputs": [],
      "source": [
        "# Copyright 2021 The TensorFlow Hub Authors. All Rights Reserved.\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     http://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License.\n",
        "# =============================================================================="
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7t7KGfIwHaXz"
      },
      "source": [
        "# Image Classification with TensorFlow Hub\n",
        "\n",
        "In this colab, you'll try multiple image classification models from TensorFlow Hub and decide which one is best for your use case.\n",
        "\n",
        "Because TF Hub encourages a [consistent input convention](https://www.tensorflow.org/hub/common_saved_model_apis/images#image_input) for models that operate on images, it's easy to experiment with different architectures to find the one that best fits your needs."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MfBg1C5NB3X0"
      },
      "source": [
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/hub/tutorials/image_classification\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" /\u003eView on TensorFlow.org\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/hub/blob/master/examples/colab/image_classification.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/hub/blob/master/examples/colab/image_classification.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca href=\"https://storage.googleapis.com/tensorflow_docs/hub/examples/colab/image_classification.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/download_logo_32px.png\" /\u003eDownload notebook\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca href=\"https://tfhub.dev/google/collections/image/1\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/hub_logo_32px.png\" /\u003eSee TF Hub models\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "N8H5ufxkc2mk"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "import tensorflow_hub as hub\n",
        "\n",
        "import requests\n",
        "from PIL import Image\n",
        "from io import BytesIO\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "oKvj6lY6kZx8"
      },
      "outputs": [],
      "source": [
        "#@title Helper functions for loading image (hidden)\n",
        "\n",
        "original_image_cache = {}\n",
        "\n",
        "def preprocess_image(image):\n",
        "  image = np.array(image)\n",
        "  # reshape into shape [batch_size, height, width, num_channels]\n",
        "  img_reshaped = tf.reshape(image, [1, image.shape[0], image.shape[1], image.shape[2]])\n",
        "  # Use `convert_image_dtype` to convert to floats in the [0,1] range.\n",
        "  image = tf.image.convert_image_dtype(img_reshaped, tf.float32)\n",
        "  return image\n",
        "\n",
        "def load_image_from_url(url):\n",
        "  \"\"\"Returns an image with shape [1, height, width, num_channels].\"\"\"\n",
        "  response = requests.get(url)\n",
        "  image = Image.open(BytesIO(response.content))\n",
        "  image = preprocess_image(image)\n",
        "  return image\n",
        "\n",
        "def load_image(image_url, image_size=256, dynamic_size=False, max_dynamic_size=512):\n",
        "  \"\"\"Loads and preprocesses images.\"\"\"\n",
        "  # Cache image file locally.\n",
        "  if image_url in original_image_cache:\n",
        "    img = original_image_cache[image_url]\n",
        "  elif image_url.startswith('https://'):\n",
        "    img = load_image_from_url(image_url)\n",
        "  else:\n",
        "    fd = tf.io.gfile.GFile(image_url, 'rb')\n",
        "    img = preprocess_image(Image.open(fd))\n",
        "  original_image_cache[image_url] = img\n",
        "  # Load and convert to float32 numpy array, add batch dimension, and normalize to range [0, 1].\n",
        "  img_raw = img\n",
        "  if tf.reduce_max(img) \u003e 1.0:\n",
        "    img = img / 255.\n",
        "  if len(img.shape) == 3:\n",
        "    img = tf.stack([img, img, img], axis=-1)\n",
        "  if not dynamic_size:\n",
        "    img = tf.image.resize_with_pad(img, image_size, image_size)\n",
        "  elif img.shape[1] \u003e max_dynamic_size or img.shape[2] \u003e max_dynamic_size:\n",
        "    img = tf.image.resize_with_pad(img, max_dynamic_size, max_dynamic_size)\n",
        "  return img, img_raw\n",
        "\n",
        "def show_image(image, title=''):\n",
        "  image_size = image.shape[1]\n",
        "  w = (image_size * 6) // 320\n",
        "  plt.figure(figsize=(w, w))\n",
        "  plt.imshow(image[0], aspect='equal')\n",
        "  plt.axis('off')\n",
        "  plt.title(title)\n",
        "  plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ws1AMDT_CDPq"
      },
      "source": [
        "Select an Image Classification Model. After that, some internal variables are set and the labels file is downloaded and prepared for use.\n",
        "\n",
        "There are some technical differences between the models, like different input size, model size, accuracy, and inference time. Here you can change the model you are using until you find the one most suitable for your use case.\n",
        "\n",
        "The handle (url) of the model is printed for your convenience. More documentation about each model is available there.\n",
        "\n",
        "Note: All these models were trained on the ImageNet dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "iQ3aamrBfs-c"
      },
      "outputs": [],
      "source": [
        "#@title Select an Image Classification model\n",
        "\n",
        "image_size = 224\n",
        "dynamic_size = False\n",
        "\n",
        "model_name = \"mobilenet_v3_small_100_224\" # @param ['bit_s-r50x1', 'efficientnet_b0', 'efficientnet_b1', 'efficientnet_b2', 'efficientnet_b3', 'efficientnet_b4', 'efficientnet_b5', 'efficientnet_b6', 'efficientnet_b7', 'inception_v3', 'inception_resnet_v2', 'mobilenet_v2_100_224', 'mobilenet_v2_130_224', 'mobilenet_v2_140_224', 'mobilenet_v3_large_100_224', 'mobilenet_v3_large_075_224', 'mobilenet_v3_small_100_224', 'mobilenet_v3_small_075_224', 'nasnet_large', 'nasnet_mobile', 'pnasnet_large', 'resnet_v1_50', 'resnet_v1_101', 'resnet_v1_152', 'resnet_v2_50', 'resnet_v2_101', 'resnet_v2_152']\n",
        "\n",
        "model_handle_map = {\n",
        "  \"efficientnet_b0\": \"https://tfhub.dev/tensorflow/efficientnet/b0/classification/1\",\n",
        "  \"efficientnet_b1\": \"https://tfhub.dev/tensorflow/efficientnet/b1/classification/1\",\n",
        "  \"efficientnet_b2\": \"https://tfhub.dev/tensorflow/efficientnet/b2/classification/1\",\n",
        "  \"efficientnet_b3\": \"https://tfhub.dev/tensorflow/efficientnet/b3/classification/1\",\n",
        "  \"efficientnet_b4\": \"https://tfhub.dev/tensorflow/efficientnet/b4/classification/1\",\n",
        "  \"efficientnet_b5\": \"https://tfhub.dev/tensorflow/efficientnet/b5/classification/1\",\n",
        "  \"efficientnet_b6\": \"https://tfhub.dev/tensorflow/efficientnet/b6/classification/1\",\n",
        "  \"efficientnet_b7\": \"https://tfhub.dev/tensorflow/efficientnet/b7/classification/1\",\n",
        "  \"bit_s-r50x1\": \"https://tfhub.dev/google/bit/s-r50x1/ilsvrc2012_classification/1\",\n",
        "  \"inception_v3\": \"https://tfhub.dev/google/imagenet/inception_v3/classification/4\",\n",
        "  \"inception_resnet_v2\": \"https://tfhub.dev/google/imagenet/inception_resnet_v2/classification/4\",\n",
        "  \"resnet_v1_50\": \"https://tfhub.dev/google/imagenet/resnet_v1_50/classification/4\",\n",
        "  \"resnet_v1_101\": \"https://tfhub.dev/google/imagenet/resnet_v1_101/classification/4\",\n",
        "  \"resnet_v1_152\": \"https://tfhub.dev/google/imagenet/resnet_v1_152/classification/4\",\n",
        "  \"resnet_v2_50\": \"https://tfhub.dev/google/imagenet/resnet_v2_50/classification/4\",\n",
        "  \"resnet_v2_101\": \"https://tfhub.dev/google/imagenet/resnet_v2_101/classification/4\",\n",
        "  \"resnet_v2_152\": \"https://tfhub.dev/google/imagenet/resnet_v2_152/classification/4\",\n",
        "  \"nasnet_large\": \"https://tfhub.dev/google/imagenet/nasnet_large/classification/4\",\n",
        "  \"nasnet_mobile\": \"https://tfhub.dev/google/imagenet/nasnet_mobile/classification/4\",\n",
        "  \"pnasnet_large\": \"https://tfhub.dev/google/imagenet/pnasnet_large/classification/4\",\n",
        "  \"mobilenet_v2_100_224\": \"https://tfhub.dev/google/imagenet/mobilenet_v2_100_224/classification/4\",\n",
        "  \"mobilenet_v2_130_224\": \"https://tfhub.dev/google/imagenet/mobilenet_v2_130_224/classification/4\",\n",
        "  \"mobilenet_v2_140_224\": \"https://tfhub.dev/google/imagenet/mobilenet_v2_140_224/classification/4\",\n",
        "  \"mobilenet_v3_small_100_224\": \"https://tfhub.dev/google/imagenet/mobilenet_v3_small_100_224/classification/5\",\n",
        "  \"mobilenet_v3_small_075_224\": \"https://tfhub.dev/google/imagenet/mobilenet_v3_small_075_224/classification/5\",\n",
        "  \"mobilenet_v3_large_100_224\": \"https://tfhub.dev/google/imagenet/mobilenet_v3_large_100_224/classification/5\",\n",
        "  \"mobilenet_v3_large_075_224\": \"https://tfhub.dev/google/imagenet/mobilenet_v3_large_075_224/classification/5\",\n",
        "}\n",
        "\n",
        "model_image_size_map = {\n",
        "  \"efficientnet_b0\": 224,\n",
        "  \"efficientnet_b1\": 240,\n",
        "  \"efficientnet_b2\": 260,\n",
        "  \"efficientnet_b3\": 300,\n",
        "  \"efficientnet_b4\": 380,\n",
        "  \"efficientnet_b5\": 456,\n",
        "  \"efficientnet_b6\": 528,\n",
        "  \"efficientnet_b7\": 600,\n",
        "  \"inception_v3\": 299,\n",
        "  \"inception_resnet_v2\": 299,\n",
        "  \"mobilenet_v2_100_224\": 224,\n",
        "  \"mobilenet_v2_130_224\": 224,\n",
        "  \"mobilenet_v2_140_224\": 224,\n",
        "  \"nasnet_large\": 331,\n",
        "  \"nasnet_mobile\": 224,\n",
        "  \"pnasnet_large\": 331,\n",
        "  \"resnet_v1_50\": 224,\n",
        "  \"resnet_v1_101\": 224,\n",
        "  \"resnet_v1_152\": 224,\n",
        "  \"resnet_v2_50\": 224,\n",
        "  \"resnet_v2_101\": 224,\n",
        "  \"resnet_v2_152\": 224,\n",
        "  \"mobilenet_v3_small_100_224\": 224,\n",
        "  \"mobilenet_v3_small_075_224\": 224,\n",
        "  \"mobilenet_v3_large_100_224\": 224,\n",
        "  \"mobilenet_v3_large_075_224\": 224,\n",
        "}\n",
        "\n",
        "model_handle = model_handle_map[model_name]\n",
        "\n",
        "print(f\"Selected model: {model_name} : {model_handle}\")\n",
        "\n",
        "\n",
        "max_dynamic_size = 512\n",
        "if model_name in model_image_size_map:\n",
        "  image_size = model_image_size_map[model_name]\n",
        "  dynamic_size = False\n",
        "  print(f\"Images will be converted to {image_size}x{image_size}\")\n",
        "else:\n",
        "  dynamic_size = True\n",
        "  print(f\"Images will be capped to a max size of {max_dynamic_size}x{max_dynamic_size}\")\n",
        "\n",
        "labels_file = \"https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt\"\n",
        "\n",
        "#download labels and creates a maps\n",
        "downloaded_file = tf.keras.utils.get_file(\"labels.txt\", origin=labels_file)\n",
        "\n",
        "classes = []\n",
        "i = 0\n",
        "with open(downloaded_file) as f:\n",
        "  labels = f.readlines()\n",
        "  classes = [l.strip() for l in labels[1:]]\n",
        "  i += 1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vxcASidjBAE8"
      },
      "source": [
        "You can select one of the images below, or use your own image. Just remember that the input sizes for the models vary and some of them use a dynamic input size (enabling inference on the unscaled image). Given that, the method `load_image` will already rescale the image to the expected format."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "o2rMsr4CgET2"
      },
      "outputs": [],
      "source": [
        "#@title Select an Input Image\n",
        "\n",
        "image_name = \"turtle\" # @param ['tiger', 'bus', 'car', 'cat', 'dog', 'apple', 'banana', 'turtle', 'flamingo', 'piano', 'honeycomb', 'teapot']\n",
        "\n",
        "images_for_test_map = {\n",
        "    \"tiger\": \"https://upload.wikimedia.org/wikipedia/commons/b/b0/Bengal_tiger_%28Panthera_tigris_tigris%29_female_3_crop.jpg\",\n",
        "    #by Charles James Sharp, CC BY-SA 4.0 \u003chttps://creativecommons.org/licenses/by-sa/4.0\u003e, via Wikimedia Commons\n",
        "    \"bus\": \"https://upload.wikimedia.org/wikipedia/commons/6/63/LT_471_%28LTZ_1471%29_Arriva_London_New_Routemaster_%2819522859218%29.jpg\",\n",
        "    #by Martin49 from London, England, CC BY 2.0 \u003chttps://creativecommons.org/licenses/by/2.0\u003e, via Wikimedia Commons\n",
        "    \"car\": \"https://upload.wikimedia.org/wikipedia/commons/4/49/2013-2016_Toyota_Corolla_%28ZRE172R%29_SX_sedan_%282018-09-17%29_01.jpg\",\n",
        "    #by EurovisionNim, CC BY-SA 4.0 \u003chttps://creativecommons.org/licenses/by-sa/4.0\u003e, via Wikimedia Commons\n",
        "    \"cat\": \"https://upload.wikimedia.org/wikipedia/commons/4/4d/Cat_November_2010-1a.jpg\",\n",
        "    #by Alvesgaspar, CC BY-SA 3.0 \u003chttps://creativecommons.org/licenses/by-sa/3.0\u003e, via Wikimedia Commons\n",
        "    \"dog\": \"https://upload.wikimedia.org/wikipedia/commons/archive/a/a9/20090914031557%21Saluki_dog_breed.jpg\",\n",
        "    #by Craig Pemberton, CC BY-SA 3.0 \u003chttps://creativecommons.org/licenses/by-sa/3.0\u003e, via Wikimedia Commons\n",
        "    \"apple\": \"https://upload.wikimedia.org/wikipedia/commons/1/15/Red_Apple.jpg\",\n",
        "    #by Abhijit Tembhekar from Mumbai, India, CC BY 2.0 \u003chttps://creativecommons.org/licenses/by/2.0\u003e, via Wikimedia Commons\n",
        "    \"banana\": \"https://upload.wikimedia.org/wikipedia/commons/1/1c/Bananas_white_background.jpg\",\n",
        "    #by fir0002  flagstaffotos [at] gmail.com\t\tCanon 20D + Tamron 28-75mm f/2.8, GFDL 1.2 \u003chttp://www.gnu.org/licenses/old-licenses/fdl-1.2.html\u003e, via Wikimedia Commons\n",
        "    \"turtle\": \"https://upload.wikimedia.org/wikipedia/commons/8/80/Turtle_golfina_escobilla_oaxaca_mexico_claudio_giovenzana_2010.jpg\",\n",
        "    #by Claudio Giovenzana, CC BY-SA 3.0 \u003chttps://creativecommons.org/licenses/by-sa/3.0\u003e, via Wikimedia Commons\n",
        "    \"flamingo\": \"https://upload.wikimedia.org/wikipedia/commons/b/b8/James_Flamingos_MC.jpg\",\n",
        "    #by Christian Mehlführer, User:Chmehl, CC BY 3.0 \u003chttps://creativecommons.org/licenses/by/3.0\u003e, via Wikimedia Commons\n",
        "    \"piano\": \"https://upload.wikimedia.org/wikipedia/commons/d/da/Steinway_%26_Sons_upright_piano%2C_model_K-132%2C_manufactured_at_Steinway%27s_factory_in_Hamburg%2C_Germany.png\",\n",
        "    #by \"Photo: © Copyright Steinway \u0026 Sons\", CC BY-SA 3.0 \u003chttps://creativecommons.org/licenses/by-sa/3.0\u003e, via Wikimedia Commons\n",
        "    \"honeycomb\": \"https://upload.wikimedia.org/wikipedia/commons/f/f7/Honey_comb.jpg\",\n",
        "    #by Merdal, CC BY-SA 3.0 \u003chttp://creativecommons.org/licenses/by-sa/3.0/\u003e, via Wikimedia Commons\n",
        "    \"teapot\": \"https://upload.wikimedia.org/wikipedia/commons/4/44/Black_tea_pot_cropped.jpg\",\n",
        "    #by Mendhak, CC BY-SA 2.0 \u003chttps://creativecommons.org/licenses/by-sa/2.0\u003e, via Wikimedia Commons\n",
        "}\n",
        "\n",
        "img_url = images_for_test_map[image_name]\n",
        "image, original_image = load_image(img_url, image_size, dynamic_size, max_dynamic_size)\n",
        "show_image(image, 'Scaled image')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CMwWx8_8Aw3X"
      },
      "source": [
        "Now that the model was chosen, loading it with TensorFlow Hub is simple.\n",
        "\n",
        "This also calls the model with a random input as a \"warmup\" run. Subsequent calls are often much faster, and you can compare this with the latency below.\n",
        "\n",
        "*Note:* models that use a dynamic size might need a fresh \"warmup\" run for each image size."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LRAccT3UhRga"
      },
      "outputs": [],
      "source": [
        "classifier = hub.load(model_handle)\n",
        "\n",
        "input_shape = image.shape\n",
        "warmup_input = tf.random.uniform(input_shape, 0, 1.0)\n",
        "%time warmup_logits = classifier(warmup_input).numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e7vkdUqpBkfE"
      },
      "source": [
        "Everything is ready for inference. Here you can see the top 5 results from the model for the selected image."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "I0QNHg3bk-G1"
      },
      "outputs": [],
      "source": [
        "# Run model on image\n",
        "%time probabilities = tf.nn.softmax(classifier(image)).numpy()\n",
        "\n",
        "top_5 = tf.argsort(probabilities, axis=-1, direction=\"DESCENDING\")[0][:5].numpy()\n",
        "np_classes = np.array(classes)\n",
        "\n",
        "# Some models include an additional 'background' class in the predictions, so\n",
        "# we must account for this when reading the class labels.\n",
        "includes_background_class = probabilities.shape[1] == 1001\n",
        "\n",
        "for i, item in enumerate(top_5):\n",
        "  class_index = item if not includes_background_class else item - 1\n",
        "  line = f'({i+1}) {class_index:4} - {classes[class_index]}: {probabilities[0][top_5][i]}'\n",
        "  print(line)\n",
        "\n",
        "show_image(image, '')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e4IJrq5eZDWl"
      },
      "source": [
        "## Learn More\n",
        "\n",
        "If you want to learn more and try how to do Transfer Learning with these models you can try this tutorial: [Transfer Learning for Image classification](https://www.tensorflow.org/hub/tutorials/tf2_image_retraining) \n",
        "\n",
        "If you want to check on more image models you can check them out on [tfhub.dev](https://tfhub.dev/s?module-type=image-augmentation,image-classification,image-classification-logits,image-classifier,image-feature-vector,image-generator,image-object-detection,image-others,image-pose-detection,image-segmentation,image-style-transfer,image-super-resolution,image-rnn-agent)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "Image Classification with TF Hub",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
